import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import basic_func.general_func as bf
from matplotlib.colors import Normalize

def cmap_data(df,params,ax):
    # 创建颜色映射
    norm = Normalize(vmin=df.iloc[:, 0].min(), vmax=df.iloc[:, 0].max())  # 归一化数据
    cmap = plt.colormaps[params['cmap_color']]  # 颜色映射
    colors = cmap(norm(df.iloc[:, 0].values))  # 生成 RGBA 颜色数组 (N, 4)
    ax.scatter(df.index,df.iloc[:,0],
               s=params['marker_size']**2,
               facecolor=colors,
               edgecolor = 'k',
               zorder=3)

def plot_stem(params):
    fig,ax = plt.subplots()
    markerline,stemlines,baseline = ax.stem(df.index,
                                             df.iloc[:,0],
                                             bottom = params['base_bottom'])
    if type(params['cmap_color'])==str:
        cmap_data(df,params,ax)
    else:
        plt.setp(markerline,marker=params['marker'],
                 markersize=params['marker_size'],
                 color=params['marker_color'])
    plt.setp(stemlines,linestyle=params['stem_ls'],
             color = params['stem_color'],
             lw = params['stem_lw'],
             alpha = params['stem_alpha'])
    plt.setp(baseline,color = params['base_color'],
             lw=params['base_lw'])

    plt.title(params['title'] if type(params['title'])==str else None,
              fontsize = params['title_size'],
              fontname = params['title_fname'])
    plt.show()
    
if __name__ == "__main__":
    df,params = bf.read_data('test_data/棒棒糖图.xlsx')
    bf.zh_font(params)
    plot_stem(params)
